package client

import (
	"chinaport-proxy/protocol"
	"chinaport-proxy/utils"
	"context"
	"encoding/json"
	"fmt"
	"log/slog"
	"math"
	"sync"
	"sync/atomic"
	"time"

	"github.com/gorilla/websocket"
)

// globalIDCounter is a global atomic counter for generating unique _id values
var globalIDCounter atomic.Int64

// nextID generates the next global ID, resetting to 1 if it would exceed max int32
func nextID() int {
	for {
		current := globalIDCounter.Load()
		// Reset to 1 if we're about to exceed max int32
		if current >= math.MaxInt32 {
			if globalIDCounter.CompareAndSwap(current, 1) {
				return 1
			}
			continue
		}
		// Normal increment
		next := current + 1
		if globalIDCounter.CompareAndSwap(current, next) {
			return int(next)
		}
	}
}

// pendingRequest represents a request waiting for response
type pendingRequest struct {
	id       int
	respChan chan *protocol.Response
	errChan  chan error
	timer    *time.Timer
}

// UKeyClient maintains a persistent WebSocket connection with request tracking
type UKeyClient struct {
	url      string
	timeout  time.Duration
	password string

	// WebSocket connection
	conn   *websocket.Conn
	connMu sync.RWMutex

	// Request tracking
	pending   map[int]*pendingRequest
	pendingMu sync.RWMutex

	// Connection state
	connected   atomic.Bool
	connectedCh chan struct{}
	connectedMu sync.Mutex

	// Reconnection
	reconnectCh chan struct{}
	backoff     time.Duration

	// Health check
	healthCheckInterval time.Duration
	healthCheckEnabled  bool

	// Webhook service for notifications
	webhookService interface {
		NotifyConnectionFailure(ctx context.Context, err error)
		NotifyHealthCheckFailure(ctx context.Context, err error)
	}

	// Lifecycle
	ctx    context.Context
	cancel context.CancelFunc
	wg     sync.WaitGroup
}

// NewUKeyClient creates a new UKey client with persistent connection
func NewUKeyClient(url string, timeout time.Duration, password string) *UKeyClient {
	ctx, cancel := context.WithCancel(context.Background())

	client := &UKeyClient{
		url:                 url,
		timeout:             timeout,
		password:            password,
		pending:             make(map[int]*pendingRequest),
		connectedCh:         make(chan struct{}),
		reconnectCh:         make(chan struct{}, 1),
		backoff:             1 * time.Second,
		healthCheckInterval: 30 * time.Second,
		healthCheckEnabled:  true,
		ctx:                 ctx,
		cancel:              cancel,
	}

	slog.Info("Creating UKey client", "url", url, "timeout", timeout)

	// Start connection manager
	client.start()

	return client
}

// SetWebhookService sets the webhook service for notifications
func (c *UKeyClient) SetWebhookService(ws interface {
	NotifyConnectionFailure(ctx context.Context, err error)
	NotifyHealthCheckFailure(ctx context.Context, err error)
}) {
	c.webhookService = ws
}

// start initializes the connection and starts background goroutines
func (c *UKeyClient) start() {
	// Start connection manager
	c.wg.Add(1)
	go c.connectionManager()

	// Start reader
	c.wg.Add(1)
	go c.readLoop()

	// Start health check
	if c.healthCheckEnabled {
		c.wg.Add(1)
		go c.healthCheckLoop()
	}

	// Trigger initial connection
	select {
	case c.reconnectCh <- struct{}{}:
	default:
	}
}

// connect establishes a new WebSocket connection
func (c *UKeyClient) connect() error {
	c.connMu.Lock()
	defer c.connMu.Unlock()

	// Close existing connection if any
	if c.conn != nil {
		c.conn.Close()
		c.conn = nil
	}

	slog.Debug("Attempting WebSocket connection", "url", c.url)

	// Dial new connection
	conn, _, err := websocket.DefaultDialer.Dial(c.url, nil)
	if err != nil {
		slog.Warn("Failed to connect", "error", err)
		// Notify webhook service about connection failure
		if c.webhookService != nil {
			go c.webhookService.NotifyConnectionFailure(context.Background(), err)
		}
		return err
	}

	c.conn = conn
	c.backoff = 1 * time.Second // Reset backoff on successful connection

	// Mark as connected
	if !c.connected.Load() {
		c.connected.Store(true)
		slog.Info("WebSocket connected")

		// Signal connection established
		c.connectedMu.Lock()
		if c.connectedCh != nil {
			close(c.connectedCh)
			c.connectedCh = make(chan struct{})
		}
		c.connectedMu.Unlock()
	}

	return nil
}

// connectionManager handles connection and reconnection
func (c *UKeyClient) connectionManager() {
	defer c.wg.Done()

	for {
		select {
		case <-c.ctx.Done():
			return
		case <-c.reconnectCh:
			// Attempt connection with exponential backoff
			for {
				if c.ctx.Err() != nil {
					return
				}

				if err := c.connect(); err != nil {
					slog.Info("Reconnecting", "backoff", c.backoff)

					select {
					case <-time.After(c.backoff):
						// Exponential backoff with max of 8 seconds
						c.backoff = c.backoff * 2
						if c.backoff > 8*time.Second {
							c.backoff = 8 * time.Second
						}
					case <-c.ctx.Done():
						return
					}
				} else {
					break // Successfully connected
				}
			}
		}
	}
}

// readLoop continuously reads messages from the WebSocket
func (c *UKeyClient) readLoop() {
	defer c.wg.Done()

	for {
		select {
		case <-c.ctx.Done():
			return
		default:
		}

		// Get connection
		c.connMu.RLock()
		conn := c.conn
		c.connMu.RUnlock()

		if conn == nil {
			time.Sleep(100 * time.Millisecond)
			continue
		}

		// Read message
		_, message, err := conn.ReadMessage()
		if err != nil {
			if c.connected.Load() {
				slog.Info("WebSocket disconnected", "error", err)
				c.connected.Store(false)

				// Trigger reconnection
				select {
				case c.reconnectCh <- struct{}{}:
				default:
				}
			}
			time.Sleep(100 * time.Millisecond)
			continue
		}

		slog.Debug("Received WebSocket message", "raw", utils.MaskPassword(string(message)))

		// Parse response
		var resp protocol.Response
		if err := json.Unmarshal(message, &resp); err != nil {
			// Not a protocol response, might be an event message
			slog.Debug("Received non-protocol message", "message", utils.MaskPassword(string(message)))
			continue
		}

		slog.Debug("Received protocol response", "id", resp.ID, "method", resp.Method, "status", resp.Status)

		// Route to waiting request
		c.pendingMu.RLock()
		pending, exists := c.pending[resp.ID]
		c.pendingMu.RUnlock()

		if exists {
			slog.Debug("Routing response to waiting request", "id", resp.ID)
			pending.timer.Stop()

			select {
			case pending.respChan <- &resp:
			default:
				slog.Warn("Response channel full", "id", resp.ID)
			}

			// Clean up
			c.pendingMu.Lock()
			delete(c.pending, resp.ID)
			c.pendingMu.Unlock()
		} else {
			slog.Warn("Received response for unknown request", "id", resp.ID)
		}
	}
}

// healthCheckLoop performs periodic health checks
func (c *UKeyClient) healthCheckLoop() {
	defer c.wg.Done()

	ticker := time.NewTicker(c.healthCheckInterval)
	defer ticker.Stop()

	for {
		select {
		case <-c.ctx.Done():
			return
		case <-ticker.C:
			if c.connected.Load() {
				slog.Debug("Performing health check")

				// Use SpcGetCertNo as health check
				ctx, cancel := context.WithTimeout(c.ctx, 5*time.Second)
				_, err := c.CallWithContext(ctx, protocol.MethodGetCertNo, struct{}{})
				cancel()

				if err != nil {
					slog.Warn("Health check failed", "error", err)

					// Notify webhook service about health check failure
					if c.webhookService != nil {
						go c.webhookService.NotifyHealthCheckFailure(context.Background(), err)
					}

					// Mark as disconnected and trigger reconnection
					c.connected.Store(false)
					select {
					case c.reconnectCh <- struct{}{}:
					default:
					}
				} else {
					slog.Debug("Health check successful")
				}
			}
		}
	}
}

// CallWithRequestID sends a request with request ID for logging
func (c *UKeyClient) CallWithRequestID(method protocol.Method, args interface{}, requestID string) (*protocol.Response, error) {
	ctx, cancel := context.WithTimeout(c.ctx, c.timeout)
	defer cancel()
	return c.callWithContextAndRequestID(ctx, method, args, requestID)
}

// CallWithContext sends a request with context and waits for response
func (c *UKeyClient) CallWithContext(ctx context.Context, method protocol.Method, args interface{}) (*protocol.Response, error) {
	return c.callWithContextAndRequestID(ctx, method, args, "")
}

// callWithContextAndRequestID is the internal implementation with optional request ID
func (c *UKeyClient) callWithContextAndRequestID(ctx context.Context, method protocol.Method, args interface{}, requestID string) (*protocol.Response, error) {
	// Create logger with request ID if provided
	logger := slog.Default()
	if requestID != "" {
		logger = slog.With("request_id", requestID)
	}

	// Auto-inject password into args
	if args == nil {
		args = map[string]interface{}{"passwd": c.password}
		logger.Debug("No args provided, created with password")
	} else if argsMap, ok := args.(map[string]interface{}); ok {
		// Always set the password from config (overwrite any existing password)
		argsMap["passwd"] = c.password
		logger.Debug("Added password to existing args map")
	} else {
		// Convert other types to map and add password
		args = map[string]interface{}{
			"original_args": args,
			"passwd":        c.password,
		}
		logger.Debug("Wrapped non-map args and added password")
	}

	// Check if connected, fail immediately if not
	if !c.connected.Load() {
		return nil, fmt.Errorf("WebSocket not connected")
	}

	// Generate request ID
	id := nextID()

	logger.Debug("Preparing request", "id", id, "method", method)

	// Create request
	req := protocol.NewRequest(id, method, args)
	reqData, err := req.Marshal()
	if err != nil {
		return nil, fmt.Errorf("failed to marshal request: %w", err)
	}

	logger.Debug("Sending WebSocket request", "id", id, "raw", utils.MaskPassword(string(reqData)))

	// Set up response channel
	pending := &pendingRequest{
		id:       id,
		respChan: make(chan *protocol.Response, 1),
		errChan:  make(chan error, 1),
	}

	// Register pending request
	c.pendingMu.Lock()
	c.pending[id] = pending
	c.pendingMu.Unlock()

	// Clean up on exit
	defer func() {
		if pending.timer != nil {
			pending.timer.Stop()
		}
		c.pendingMu.Lock()
		delete(c.pending, id)
		c.pendingMu.Unlock()
	}()

	// Set timeout
	timeoutDuration := c.timeout
	if deadline, ok := ctx.Deadline(); ok {
		if d := time.Until(deadline); d < timeoutDuration {
			timeoutDuration = d
		}
	}

	pending.timer = time.AfterFunc(timeoutDuration, func() {
		logger.Warn("Request timeout", "id", id, "timeout", timeoutDuration)
		select {
		case pending.errChan <- fmt.Errorf("request timeout after %v", timeoutDuration):
		default:
		}
	})

	// Send request
	c.connMu.RLock()
	conn := c.conn
	c.connMu.RUnlock()

	if conn == nil {
		return nil, fmt.Errorf("no connection available")
	}

	if err := conn.WriteMessage(websocket.TextMessage, reqData); err != nil {
		logger.Warn("Failed to send WebSocket request", "id", id, "error", err)

		// Mark as disconnected and trigger reconnection
		c.connected.Store(false)
		select {
		case c.reconnectCh <- struct{}{}:
		default:
		}

		return nil, fmt.Errorf("failed to send request: %w", err)
	}

	// Wait for response
	logger.Debug("Waiting for response", "id", id)

	select {
	case resp := <-pending.respChan:
		logger.Debug("Received response", "id", id, "status", resp.Status)
		return resp, nil
	case err := <-pending.errChan:
		logger.Warn("Request failed", "id", id, "error", err)
		return nil, err
	case <-ctx.Done():
		logger.Debug("Request cancelled", "id", id)
		return nil, ctx.Err()
	}
}

// Call sends a request and waits for response
func (c *UKeyClient) Call(method protocol.Method, args interface{}) (*protocol.Response, error) {
	ctx, cancel := context.WithTimeout(c.ctx, c.timeout)
	defer cancel()
	return c.callWithContextAndRequestID(ctx, method, args, "")
}

// Close shuts down the client
func (c *UKeyClient) Close() error {
	slog.Info("Shutting down UKey client")

	// Cancel context to stop all goroutines
	c.cancel()

	// Close WebSocket connection
	c.connMu.Lock()
	if c.conn != nil {
		c.conn.Close()
	}
	c.connMu.Unlock()

	// Wait for goroutines to finish
	c.wg.Wait()

	// Clean up pending requests
	c.pendingMu.Lock()
	pendingCount := len(c.pending)
	for _, pending := range c.pending {
		if pending.timer != nil {
			pending.timer.Stop()
		}
		select {
		case pending.errChan <- fmt.Errorf("client closed"):
		default:
		}
	}
	c.pending = make(map[int]*pendingRequest)
	c.pendingMu.Unlock()

	if pendingCount > 0 {
		slog.Warn("Cancelled pending requests", "count", pendingCount)
	}

	return nil
}
